#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
完整测试流程脚本
按正确顺序引导用户完成RKNN模型的验证和调试
"""

import os
import sys
import subprocess
import platform

def run_command(cmd, description):
    """运行命令并显示结果"""
    print(f"\n🔄 {description}")
    print(f"执行命令: {cmd}")
    print("-" * 40)
    
    try:
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        if result.returncode == 0:
            print("✅ 执行成功")
            if result.stdout:
                print("输出:")
                print(result.stdout)
        else:
            print("❌ 执行失败")
            if result.stderr:
                print("错误信息:")
                print(result.stderr)
        return result.returncode == 0
    except Exception as e:
        print(f"❌ 命令执行异常: {e}")
        return False

def check_file_exists(filepath, description):
    """检查文件是否存在"""
    if os.path.exists(filepath):
        print(f"✅ {description}: {filepath}")
        return True
    else:
        print(f"❌ {description}不存在: {filepath}")
        return False

def check_dependencies():
    """检查依赖项"""
    print("🔍 检查依赖项...")
    
    # 检查Python包
    required_packages = ['onnxruntime', 'opencv-python', 'numpy']
    missing_packages = []
    
    for package in required_packages:
        try:
            __import__(package.replace('-', '_'))
            print(f"✅ {package}: 已安装")
        except ImportError:
            print(f"❌ {package}: 未安装")
            missing_packages.append(package)
    
    if missing_packages:
        print(f"\n⚠️  缺少依赖包: {missing_packages}")
        print("安装命令:")
        print(f"pip install {' '.join(missing_packages)}")
        return False
    
    return True

def step1_check_files():
    """步骤1: 检查必要文件"""
    print("\n" + "="*50)
    print("📋 步骤1: 检查必要文件")
    print("="*50)
    
    files_to_check = [
        ("best.onnx", "ONNX模型文件"),
        ("src/rkpt.cc", "C++推理代码"),
        ("src/postprocess.cpp", "后处理代码"),
        ("include/det/postprocess.h", "后处理头文件"),
        ("CMakeLists.txt", "CMake配置文件"),
        ("build.sh", "编译脚本")
    ]
    
    all_exist = True
    for filepath, description in files_to_check:
        if not check_file_exists(filepath, description):
            all_exist = False
    
    return all_exist

def step2_test_onnx():
    """步骤2: 测试ONNX模型"""
    print("\n" + "="*50) 
    print("📋 步骤2: 验证ONNX模型")
    print("="*50)
    
    if not os.path.exists("test_onnx_model.py"):
        print("❌ test_onnx_model.py 不存在")
        return False
    
    return run_command("python test_onnx_model.py", "验证ONNX模型")

def step3_convert_rknn():
    """步骤3: 转换RKNN模型"""
    print("\n" + "="*50)
    print("📋 步骤3: 转换RKNN模型") 
    print("="*50)
    
    if not os.path.exists("onnx_to_rknn_640x640.py"):
        print("❌ onnx_to_rknn_640x640.py 不存在")
        return False
    
    return run_command("python onnx_to_rknn_640x640.py", "转换RKNN模型")

def step4_check_rknn():
    """步骤4: 检查RKNN模型信息"""
    print("\n" + "="*50)
    print("📋 步骤4: 检查RKNN模型信息")
    print("="*50)
    
    if not os.path.exists("check_rknn_info.py"):
        print("❌ check_rknn_info.py 不存在")
        return False
    
    return run_command("python check_rknn_info.py", "检查RKNN模型信息")

def step5_compile_cpp():
    """步骤5: 编译C++代码"""
    print("\n" + "="*50)
    print("📋 步骤5: 编译C++代码")
    print("="*50)
    
    # 检查是否在Linux/RK3588环境
    system = platform.system()
    if system != "Linux":
        print(f"⚠️  当前系统: {system}")
        print("C++编译需要在Linux/RK3588环境进行")
        print("请在目标硬件上执行: ./build.sh")
        return True  # 不算失败，只是需要在不同环境执行
    
    return run_command("./build.sh", "编译C++代码")

def step6_run_test():
    """步骤6: 运行测试"""
    print("\n" + "="*50)
    print("📋 步骤6: 运行检测测试")
    print("="*50)
    
    system = platform.system()
    if system != "Linux":
        print(f"⚠️  当前系统: {system}")
        print("RKNN推理需要在RK3588硬件上运行")
        print("请在目标硬件上执行: ./build_and_test.sh")
        return True
    
    return run_command("./build_and_test.sh", "运行检测测试")

def main():
    """主函数"""
    print("🚀 RKNN物资检测完整测试流程")
    print("="*50)
    print("📋 测试流程：")
    print("1. 检查必要文件")
    print("2. 验证ONNX模型（PC端）")
    print("3. 转换RKNN模型")
    print("4. 检查RKNN模型信息")
    print("5. 编译C++代码（需要RK3588环境）")
    print("6. 运行检测测试（需要RK3588环境）")
    print("="*50)
    
    # 检查依赖
    if not check_dependencies():
        print("\n❌ 依赖检查失败，请先安装缺少的包")
        return
    
    # 执行测试步骤
    steps = [
        ("检查文件", step1_check_files),
        ("ONNX验证", step2_test_onnx),
        ("RKNN转换", step3_convert_rknn), 
        ("RKNN检查", step4_check_rknn),
        ("C++编译", step5_compile_cpp),
        ("运行测试", step6_run_test)
    ]
    
    success_count = 0
    for step_name, step_func in steps:
        try:
            if step_func():
                success_count += 1
                print(f"✅ {step_name} - 完成")
            else:
                print(f"❌ {step_name} - 失败")
                
                # 对于前4步失败，询问是否继续
                if success_count < 4:
                    choice = input(f"\n⚠️  {step_name}失败，是否继续下一步？(y/n): ")
                    if choice.lower() != 'y':
                        print("测试流程中断")
                        break
        except KeyboardInterrupt:
            print(f"\n⚠️  用户中断测试流程")
            break
        except Exception as e:
            print(f"❌ {step_name}执行异常: {e}")
    
    # 结果总结
    print("\n" + "="*50)
    print("📊 测试结果总结")
    print("="*50)
    print(f"完成步骤: {success_count}/{len(steps)}")
    
    if success_count >= 4:
        print("✅ 主要验证步骤完成，可以在RK3588上进行最终测试")
        print("\n🎯 下一步操作（在RK3588硬件上）：")
        print("1. 传输所有文件到RK3588")
        print("2. 运行: ./build.sh")  
        print("3. 运行: ./build_and_test.sh")
        print("4. 观察检测结果和调试输出")
    elif success_count >= 2:
        print("⚠️  部分验证完成，建议修复失败的步骤后再继续")
    else:
        print("❌ 基础验证失败，请检查模型文件和环境配置")
    
    print("\n💡 故障排除建议：")
    print("- 如果ONNX验证失败：检查模型文件和训练配置")
    print("- 如果RKNN转换失败：检查RKNN工具链安装")  
    print("- 如果编译失败：检查RK3588开发环境配置")
    print("- 如果检测失败：调整置信度阈值或检查预处理")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n⚠️  测试流程被用户中断")
    except Exception as e:
        print(f"\n❌ 程序异常: {e}") 